package com.huaxin.hxmodulewebsocket.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class MyWebSocketHandler implements WebSocketHandler {

    // 存储所有连接的客户端 session，key 可以是用户ID、用户名或 sessionId
    private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    private final ObjectMapper objectMapper = new ObjectMapper();

    // 当客户端连接时触发
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String clientId = getClientId(session); // 可以从 URI 参数、header 或 session 中获取
        sessions.put(clientId, session);
        System.out.println("客户端 " + clientId + " 已连接，当前在线: " + sessions.size());

        // 可选：发送欢迎消息
        sendMessageToClient(session, "欢迎，" + clientId + "！");
    }

    // 当收到客户端消息时触发
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String payload = ((TextMessage) message).getPayload();
        System.out.println("收到来自 " + getClientId(session) + " 的消息: " + payload);

        // 解析消息（示例：假设消息是 JSON {"type": "broadcast|private", "to": "id", "content": "..."})
        try {
            Message msg = objectMapper.readValue(payload, Message.class);
            switch (msg.getType()) {
                case "broadcast":
                    broadcastMessage(msg.getContent());
                    break;
                case "private":
                    sendToClient(msg.getTo(), msg.getContent());
                    break;
                default:
                    sendMessageToClient(session, "未知消息类型");
            }
        } catch (Exception e) {
            sendMessageToClient(session, "消息格式错误: " + e.getMessage());
        }
    }

    // 发送消息给单个客户端
    public void sendMessageToClient(WebSocketSession session, String message) throws IOException {
        if (session.isOpen()) {
            session.sendMessage(new TextMessage(message));
        }
    }

    // 广播消息给所有客户端
    public void broadcastMessage(String message) {
        sessions.values().forEach(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(message));
                }
            } catch (IOException e) {
                System.err.println("广播消息失败: " + e.getMessage());
            }
        });
    }

    // 发送消息给指定客户端
    public void sendToClient(String clientId, String message) {
        WebSocketSession session = sessions.get(clientId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                System.err.println("发送给 " + clientId + " 失败: " + e.getMessage());
            }
        } else {
            System.out.println("客户端 " + clientId + " 不在线或已断开");
        }
    }

    // 从 session 中提取客户端标识（示例：从 URI 获取 ?id=xxx）
    private String getClientId(WebSocketSession session) {
        // 示例：ws://localhost:8080/ws?id=user123
        return session.getUri().getQuery().replace("id=", "");
        // 你也可以使用 session.getId() 作为唯一标识
        // return session.getId();
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("传输错误: " + exception.getMessage());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String clientId = getClientId(session);
        sessions.remove(clientId);
        System.out.println("客户端 " + clientId + " 已断开，当前在线: " + sessions.size());
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    // 消息对象
    public static class Message {
        private String type;      // "broadcast" 或 "private"
        private String to;        // 接收者ID（私聊时使用）
        private String content;   // 消息内容

        // Getters and Setters
        public String getType() {
            return type;
        }

        public void setType(String type) {
            this.type = type;
        }

        public String getTo() {
            return to;
        }

        public void setTo(String to) {
            this.to = to;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }
}